// Copyright 2018 ETH Zürich, Thomas Schöps
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
//    this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors
//    may be used to endorse or promote products derived from this software
//    without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.


#pragma once

#include <fstream>

#include <glog/logging.h>

#include "libvis/eigen.h"
#include "libvis/image.h"
#include "libvis/libvis.h"
#include "libvis/sophus.h"

namespace vis {

// Point type storing only a position attribute.
template <typename T>
struct Point {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  
  inline Point(const T& position)
      : position_(position) {}
  
  inline T& position() { return position_; }
  inline const T& position() const { return position_; }
  
 private:
  T position_;
};


// Point type storing position and color.
template <typename PositionT, typename ColorT>
struct PointC {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  
  inline PointC(const PositionT& position, const ColorT& color)
      : position_(position), color_(color) {}
  
  inline PositionT& position() { return position_; }
  inline const PositionT& position() const { return position_; }
  
  inline ColorT& color() { return color_; }
  inline const ColorT& color() const { return color_; }
  
 private:
  PositionT position_;
  ColorT color_;
};


// Traits type for querying the presence of point attributes.
// Must have a specialization for each point type that is used.
template <typename PointT>
struct PointTraits {};

template<typename _PositionT>
struct PointTraits<Point<_PositionT>> {
  typedef _PositionT PositionT;
  
  static const bool has_color = false;
  typedef Vec2u8 ColorT;  // The chosen default type here does not matter.
};

template<typename _PositionT, typename _ColorT>
struct PointTraits<PointC<_PositionT, _ColorT>> {
  typedef _PositionT PositionT;
  
  static const bool has_color = true;
  typedef _ColorT ColorT;
};


// Traits type for applying a cast to Eigen Matrix types or to scalar types,
// either using static_cast<>() or using Eigen::Matrix::cast<>(),
// by passing for example unsigned char or Eigen::Matrix<float, 1, 1> as T.
template <typename InputT>
struct CastEigenOrScalar {
  template <typename T>
  inline T Cast(const InputT& input) {
    return static_cast<T>(input);
  }
};

template<typename _Scalar, int _Rows, int _Cols, int _Options, int _MaxRows, int _MaxCols>
struct CastEigenOrScalar<Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>> {
  template <typename T>
  inline T Cast(const Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>& input) {
    return input.template cast<typename T::Scalar>();
  }
};


// Generic point cloud type, templated with the point type T.
// The point types should have some attributes with common names such that they
// can be used by PointCloud. At the moment, this is position() which is
// expected to return a writable reference to an Eigen::Vector compatible type,
// and if the point types has colors, color() to return the color.
template <typename PointT>
class PointCloud {
 public:
  using PositionT = typename PointTraits<PointT>::PositionT;
  
  // Creates an empty point cloud.
  PointCloud()
      : data_(nullptr), size_(0), alignment_(0) {}
  
  // Creates a deep copy of the other point cloud.
  PointCloud(const PointCloud<PointT>& other)
      : data_(nullptr), size_(0), alignment_(0) {
    Resize(other.size_, other.alignment_);
    memcpy(data_, other.data_, SizeInBytes());
  }
  
  // Creates a point cloud with the given point count. Does not initialize the
  // point memory. The memory alignment is chosen automatically.
  PointCloud(usize size)
      : data_(nullptr), size_(0), alignment_(0) {
    Resize(size);
  }
  
  // Creates a point cloud with the given point count and alignment (in bytes,
  // must be a power of two and multiple of sizeof(void*), or 1 for no
  // alignment).
  PointCloud(usize size, usize alignment)
      : data_(nullptr), size_(0), alignment_(0) {
    Resize(size, alignment);
  }
  
  ~PointCloud() {
    // Does nothing if data_ is nullptr.
    free(data_);
  }
  
  // Changes the size. Re-allocates the point buffer if the new size is
  // different from the current one. Does not preserve the point data.
  void Resize(usize size) {
    if (data_ && size_ == size) {
      return;
    }
    Resize(size, 1);
  }
  
  // Changes the size and alignment (in bytes, must be a power of two and
  // multiple of sizeof(void*), or 1 for no alignment). Re-allocates the point
  // buffer if the new settings are different from the current ones. Does not
  // preserve the point data.
  void Resize(usize size, usize alignment) {
    if (data_ && size_ == size && alignment_ == alignment) {
      return;
    }
    
    // Does nothing if data_ is nullptr.
    free(data_);
    data_ = nullptr;
    
    int return_value;
    if (alignment == 1) {
      data_ = reinterpret_cast<PointT*>(malloc(size * sizeof(PointT)));
      return_value = (data_ == nullptr) ? (-1) : 0;
    } else {
      return_value = posix_memalign(reinterpret_cast<void**>(&data_), alignment, size * sizeof(PointT));
    }
    if (return_value != 0) {
      // An error ocurred.
      if (return_value == EINVAL) {
        // The alignment argument was not a power of two, or was not a multiple of
        // sizeof(void*).
        // TODO
        LOG(FATAL) << "return_value == EINVAL";
      } else if (return_value == ENOMEM) {
        // There was insufficient memory to fulfill the allocation request.
        // TODO
        LOG(FATAL) << "return_value == ENOMEM";
      } else {
        // Unknown error.
        // TODO
        LOG(FATAL) << "Unknown error";
      }
    }
    
    size_ = size;
    alignment_ = alignment;
  }
  
  // Computes the axis-aligned bounding box extents for the point cloud. Does
  // not do anything for empty point clouds.
  // TODO: Compare performance between the current method and the one which
  //       is commented out, and additionally, to using _min and _max directly
  //       instead of using local variables on the stack.
  void ComputeMinMax(PositionT* _min, PositionT* _max) const {
    if (empty()) {
      return;
    }
    
    PositionT min = data_[0].position();
    PositionT max = data_[0].position();
    
    for (usize i = 1; i < size_; ++ i) {
      for (int c = 0; c < PositionT::RowsAtCompileTime; ++ c) {
        if (data_[i].position().coeffRef(c) <
            min.coeffRef(c)) {
          min.coeffRef(c) = data_[i].position().coeffRef(c);
        }
        if (data_[i].position().coeffRef(c) >
            max.coeffRef(c)) {
          max.coeffRef(c) = data_[i].position().coeffRef(c);
        }
      }
      
//       *min = min->cwiseMin(data_[i].position());
//       *max = max->cwiseMax(data_[i].position());
    }
    
    *_min = min;
    *_max = max;
  }
  
  // Un-projects the depth map pixels into 3D space, creates a point for each
  // valid pixel. Only the depth attribute of the points is set, possible other
  // attributes remain uninitialized.
  //
  // \sa SetFromRGBDImage
  template <typename DepthT, typename CameraT>
  void SetFromDepthImage(const Image<DepthT> depth_image, bool depth_is_inverse, DepthT invalid_depth_value, const CameraT& camera) {
    // Count the valid pixels in the depth image.
    usize point_count = 0;
    for (u32 y = 0; y < depth_image.height(); ++ y) {
      const DepthT* ptr = depth_image.row(y);
      const DepthT* end = ptr + depth_image.width();
      while (ptr < end) {
        if (*ptr != invalid_depth_value) {
          ++ point_count;
        }
        ++ ptr;
      }
    }
    
    Resize(point_count);
    
    // Create the points.
    point_count = 0;
    for (u32 y = 0; y < depth_image.height(); ++ y) {
      const DepthT* ptr = depth_image.row(y);
      for (u32 x = 0; x < depth_image.width(); ++ x) {
        if (*ptr != invalid_depth_value) {
          data_[point_count].position() =
              (depth_is_inverse ? (1.f / *ptr) : *ptr) *
              camera.UnprojectFromPixelCenterConv(Vec2f(x, y)).template cast<typename PositionT::Scalar>();
          ++ point_count;
        }
        ++ ptr;
      }
    }
  }
  
  // Un-projects the depth map pixels into 3D space, creates a point for each
  // valid pixel. This function can only be called on point clouds with point
  // types that contain a color attribute.
  //
  // \sa SetFromDepthImage
  template <typename DepthT, typename ColorT, typename CameraT>
  void SetFromRGBDImage(
      const Image<DepthT> depth_image,
      bool depth_is_inverse,
      DepthT invalid_depth_value,
      const Image<ColorT> color_image,
      const CameraT& camera) {
    // Count the valid pixels in the depth image.
    usize point_count = 0;
    for (u32 y = 0; y < depth_image.height(); ++ y) {
      const DepthT* ptr = depth_image.row(y);
      const DepthT* end = ptr + depth_image.width();
      while (ptr < end) {
        if (*ptr != invalid_depth_value) {
          ++ point_count;
        }
        ++ ptr;
      }
    }
    
    Resize(point_count);
    
    // Create the points.
    point_count = 0;
    for (u32 y = 0; y < depth_image.height(); ++ y) {
      const DepthT* d_ptr = depth_image.row(y);
      const ColorT* rgb_ptr = color_image.row(y);
      for (u32 x = 0; x < depth_image.width(); ++ x) {
        if (*d_ptr != invalid_depth_value) {
          data_[point_count].position() =
              (depth_is_inverse ? (1.f / *d_ptr) : *d_ptr) *
              camera.UnprojectFromPixelCenterConv(Vec2f(x, y)).template cast<typename PositionT::Scalar>();
          data_[point_count].color() =
              CastEigenOrScalar<ColorT>().template Cast<typename PointTraits<PointT>::ColorT>(*rgb_ptr);
          ++ point_count;
        }
        ++ d_ptr;
        ++ rgb_ptr;
      }
    }
  }
  
  // Transforms all points in the cloud by left-multiplication with the given
  // Sophus SE3 transformation.
  // 
  // There is a small overhead in computing the rotation matrix from the SE3's
  // quaternion. If transforming a large number of point clouds, it would be
  // faster to pre-compute the rotation matrix only once.
  template <typename Derived>
  void Transform(const Sophus::SE3Base<Derived>& transform) {
    // Convert the rotation quaternion to a matrix for faster point
    // multiplication.
    Matrix<typename PositionT::Scalar,
           PositionT::RowsAtCompileTime,
           PositionT::RowsAtCompileTime> rotation =
        transform.rotationMatrix().template cast<typename PositionT::Scalar>();
    Matrix<typename PositionT::Scalar,
          PositionT::RowsAtCompileTime, 1> translation =
        transform.translation().template cast<typename PositionT::Scalar>();
    
    for (usize i = 0; i < size_; ++ i) {
      data_[i].position() = rotation * data_[i].position() + translation;
    }
  }
  
  bool WriteAsOBJ(const char* path) {
    ofstream file_stream(path, std::ios::out);
    if (!file_stream) {
      return false;
    }
    
    bool result = WriteAsOBJ(&file_stream);
    
    file_stream.close();
    return result;
  }
  
  template<typename _CharT, typename _Traits>
  bool WriteAsOBJ(basic_ostream<_CharT,_Traits>* stream) {
    for (usize i = 0; i < size_; ++ i) {
      const PointT& point = data_[i];
      *stream << "v " << point.position().x()
              << " " << point.position().y()
              << " " << point.position().z();
      if (PointTraits<PointT>::has_color) {
        constexpr float kNormalizationFactor =
            1.0f / numeric_limits<typename PointTraits<PointT>::ColorT::Scalar>::max();
        *stream << " " << (kNormalizationFactor * point.color().x())
                << " " << (kNormalizationFactor * point.color().y())
                << " " << (kNormalizationFactor * point.color().z());
      }
      *stream << std::endl;
    }
    return true;
  }
  
  // Returns the i-th point in the cloud.
  inline PointT& operator[](int i) { return data_[i]; }
  inline const PointT& operator[](int i) const { return data_[i]; }
  
  inline PointT& at(int i) { return data_[i]; }
  inline const PointT& at(int i) const { return data_[i]; }
  
  // Returns a pointer to the data.
  inline const PointT* data() const { return data_; }
  inline PointT* data_mutable() { return data_; }
  
  // Returns whether the size of the point cloud is zero.
  inline bool empty() const { return size_ == 0; }
  
  // Returns the number of points in the cloud.
  inline usize size() const { return size_; }
  
  // Returns the size of the points in the cloud in bytes.
  inline usize SizeInBytes() const { return size_ * sizeof(PointT); }
  
 private:
  PointT* data_;
  usize size_;
  usize alignment_;
};


typedef Point<Vec3f> Point3f;
typedef PointC<Vec3f, u8> Point3fCu8;
typedef PointC<Vec3f, Vec3u8> Point3fC3u8;

typedef PointCloud<Point3f> Point3fCloud;
typedef PointCloud<Point3fCu8> Point3fCu8Cloud;
typedef PointCloud<Point3fC3u8> Point3fC3u8Cloud;

}
